/**
 * @file sync_bn_training_update.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "kernel_operator.h"

constexpr int32_t TOTAL_LENGTH = 2048;                            // total length of data (C)
constexpr int32_t USE_CORE_NUM = 8;                               // num of core used
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM;     // length computed of each core
constexpr int32_t TILE_NUM = 8;                                   // split data into 8 tiles for each core
constexpr int32_t BUFFER_NUM = 2;                                 // tensor num for each queue
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; // separate to 2 parts, due to double buffer

class KernelSyncBNTrainingUpdate {
public:
    __aicore__ inline KernelSyncBNTrainingUpdate() {}
    
    // 初始化函数，完成内存初始化相关操作
    __aicore__ inline void Init(GM_ADDR local_mean, GM_ADDR local_var, 
                                GM_ADDR global_mean, GM_ADDR global_var,
                                GM_ADDR updated_mean, GM_ADDR updated_var,
                                float momentum, float epsilon)
    {
        // 设置全局缓冲区
        localMeanGm.SetGlobalBuffer((__gm__ float *)local_mean + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        localVarGm.SetGlobalBuffer((__gm__ float *)local_var + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        globalMeanGm.SetGlobalBuffer((__gm__ float *)global_mean + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        globalVarGm.SetGlobalBuffer((__gm__ float *)global_var + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        updatedMeanGm.SetGlobalBuffer((__gm__ float *)updated_mean + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        updatedVarGm.SetGlobalBuffer((__gm__ float *)updated_var + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        
        // 初始化队列缓冲区
        pipe.InitBuffer(inQueueLocalMean, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(inQueueLocalVar, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(inQueueGlobalMean, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(inQueueGlobalVar, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(outQueueUpdatedMean, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(outQueueUpdatedVar, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        
        // 存储属性参数
        this->momentum = momentum;
        this->epsilon = epsilon;
        this->one_minus_momentum = 1.0f - momentum;
    }
    
    // 核心处理函数
    __aicore__ inline void Process()
    {
        int32_t loopCount = TILE_NUM * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    // 搬入函数
    __aicore__ inline void CopyIn(int32_t progress)
    {
        // 分配本地张量
        AscendC::LocalTensor<float> localMeanLocal = inQueueLocalMean.AllocTensor<float>();
        AscendC::LocalTensor<float> localVarLocal = inQueueLocalVar.AllocTensor<float>();
        AscendC::LocalTensor<float> globalMeanLocal = inQueueGlobalMean.AllocTensor<float>();
        AscendC::LocalTensor<float> globalVarLocal = inQueueGlobalVar.AllocTensor<float>();
        
        // 从全局内存拷贝数据到本地内存
        AscendC::DataCopy(localMeanLocal, localMeanGm[progress * TILE_LENGTH], TILE_LENGTH);
        AscendC::DataCopy(localVarLocal, localVarGm[progress * TILE_LENGTH], TILE_LENGTH);
        AscendC::DataCopy(globalMeanLocal, globalMeanGm[progress * TILE_LENGTH], TILE_LENGTH);
        AscendC::DataCopy(globalVarLocal, globalVarGm[progress * TILE_LENGTH], TILE_LENGTH);
        
        // 将数据放入队列
        inQueueLocalMean.EnQue(localMeanLocal);
        inQueueLocalVar.EnQue(localVarLocal);
        inQueueGlobalMean.EnQue(globalMeanLocal);
        inQueueGlobalVar.EnQue(globalVarLocal);
    }
    
    // 计算函数
    __aicore__ inline void Compute(int32_t progress)
    {
        // 从队列中取出数据
        AscendC::LocalTensor<float> localMeanLocal = inQueueLocalMean.DeQue<float>();
        AscendC::LocalTensor<float> localVarLocal = inQueueLocalVar.DeQue<float>();
        AscendC::LocalTensor<float> globalMeanLocal = inQueueGlobalMean.DeQue<float>();
        AscendC::LocalTensor<float> globalVarLocal = inQueueGlobalVar.DeQue<float>();
        
        // 分配输出张量
        AscendC::LocalTensor<float> updatedMeanLocal = outQueueUpdatedMean.AllocTensor<float>();
        AscendC::LocalTensor<float> updatedVarLocal = outQueueUpdatedVar.AllocTensor<float>();
        
        // 创建常量张量
        AscendC::LocalTensor<float> momentumTensor;
        AscendC::LocalTensor<float> oneMinusMomentumTensor;
        AscendC::LocalTensor<float> epsilonTensor;
        
        // 填充常量值
        AscendC::Duplicate(momentumTensor, momentum, TILE_LENGTH);
        AscendC::Duplicate(oneMinusMomentumTensor, one_minus_momentum, TILE_LENGTH);
        AscendC::Duplicate(epsilonTensor, epsilon, TILE_LENGTH);
        
        // 计算 updated_mean = momentum * global_mean + (1 - momentum) * local_mean
        AscendC::LocalTensor<float> temp1, temp2;
        AscendC::Mul(temp1, globalMeanLocal, momentumTensor, TILE_LENGTH);
        AscendC::Mul(temp2, localMeanLocal, oneMinusMomentumTensor, TILE_LENGTH);
        AscendC::Add(updatedMeanLocal, temp1, temp2, TILE_LENGTH);
        
        // 计算 updated_var = momentum * global_var + (1 - momentum) * local_var + epsilon
        AscendC::LocalTensor<float> temp3, temp4;
        AscendC::Mul(temp3, globalVarLocal, momentumTensor, TILE_LENGTH);
        AscendC::Mul(temp4, localVarLocal, oneMinusMomentumTensor, TILE_LENGTH);
        AscendC::Add(temp3, temp3, temp4, TILE_LENGTH);
        AscendC::Add(updatedVarLocal, temp3, epsilonTensor, TILE_LENGTH);
        
        // 将结果放入输出队列
        outQueueUpdatedMean.EnQue<float>(updatedMeanLocal);
        outQueueUpdatedVar.EnQue<float>(updatedVarLocal);
        
        // 释放输入张量
        inQueueLocalMean.FreeTensor(localMeanLocal);
        inQueueLocalVar.FreeTensor(localVarLocal);
        inQueueGlobalMean.FreeTensor(globalMeanLocal);
        inQueueGlobalVar.FreeTensor(globalVarLocal);
    }
    
    // 搬出函数
    __aicore__ inline void CopyOut(int32_t progress)
    {
        // 从输出队列中取出数据
        AscendC::LocalTensor<float> updatedMeanLocal = outQueueUpdatedMean.DeQue<float>();
        AscendC::LocalTensor<float> updatedVarLocal = outQueueUpdatedVar.DeQue<float>();
        
        // 将结果拷贝到全局内存
        AscendC::DataCopy(updatedMeanGm[progress * TILE_LENGTH], updatedMeanLocal, TILE_LENGTH);
        AscendC::DataCopy(updatedVarGm[progress * TILE_LENGTH], updatedVarLocal, TILE_LENGTH);
        
        // 释放输出张量
        outQueueUpdatedMean.FreeTensor(updatedMeanLocal);
        outQueueUpdatedVar.FreeTensor(updatedVarLocal);
    }

private:
    // 属性参数
    float momentum;
    float epsilon;
    float one_minus_momentum;
    
    // 内存管理对象
    AscendC::TPipe pipe;
    
    // 输入队列
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueLocalMean;
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueLocalVar;
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueGlobalMean;
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueGlobalVar;
    
    // 输出队列
    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueUpdatedMean;
    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueUpdatedVar;
    
    // 全局内存管理对象
    AscendC::GlobalTensor<float> localMeanGm;
    AscendC::GlobalTensor<float> localVarGm;
    AscendC::GlobalTensor<float> globalMeanGm;
    AscendC::GlobalTensor<float> globalVarGm;
    AscendC::GlobalTensor<float> updatedMeanGm;
    AscendC::GlobalTensor<float> updatedVarGm;
};

// 核函数定义
extern "C" __global__ __aicore__ void sync_bn_training_update(GM_ADDR local_mean, GM_ADDR local_var, 
                                                             GM_ADDR global_mean, GM_ADDR global_var,
                                                             GM_ADDR updated_mean, GM_ADDR updated_var,
                                                             float momentum, float epsilon)
{
    KernelSyncBNTrainingUpdate op;
    op.Init(local_mean, local_var, global_mean, global_var, 
            updated_mean, updated_var, momentum, epsilon);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void sync_bn_training_update_do(uint32_t blockDim, void *stream, 
                                uint8_t *local_mean, uint8_t *local_var,
                                uint8_t *global_mean, uint8_t *global_var,
                                uint8_t *updated_mean, uint8_t *updated_var,
                                float momentum, float epsilon)
{
    sync_bn_training_update<<<blockDim, nullptr, stream>>>(local_mean, local_var, global_mean, global_var,
                                                          updated_mean, updated_var, momentum, epsilon);
}
#endif